import torch.nn as nn
import torch.nn.init as init

from .nets.backbone import HourglassBackbone, SuperpointBackbone,SuperpointBackbone_4

from .nets.junction_decoder import SuperpointDecoder
from .nets.heatmap_decoder import PixelShuffleDecoder
from .nets.heatmap_decoder import MultiScaleAttention
from .nets.heatmap_decoder import heatmap_head
from .nets.angle_decoder import angle_decoder
from .nets.descriptor_decoder import SuperpointDescriptor


def get_model(model_cfg=None, loss_weights=None, mode="train"):
    """ Get model based on the model configuration. """
    # Check dataset config is given
    if model_cfg is None:
        raise ValueError("[Error] The model config is required!")

    # List the supported options here
    print("\n\n\t--------Initializing model----------")
    supported_arch = ["simple"]
    if not model_cfg["model_architecture"] in supported_arch:
        raise ValueError(
            "[Error] The model architecture is not in supported arch!")

    if model_cfg["model_architecture"] == "simple":
        #SOLD2Net网络模型类的实例化对象
        #出入配置python字典，初始化网络参数
        model = SOLD2Net(model_cfg)
    else:
        raise ValueError(
            "[Error] The model architecture is not in supported arch!")

    # Optionally register loss weights to the model
    if mode == "train":
        if loss_weights is not None:
            for param_name, param in loss_weights.items():
                #判断对象是否是一个已知的类型
                if isinstance(param, nn.Parameter):
                    print("\t [Debug] Adding %s with value %f to model"
                          % (param_name, param.item()))
                    #向网络添加parameter
                    model.register_parameter(param_name, param)
        else:
            raise ValueError(
                "[Error] the loss weights can not be None in dynamic weighting mode during training.")

    # Display some summary info.
    print("\tModel architecture: %s" % model_cfg["model_architecture"])     #simple
    print("\tBackbone: %s" % model_cfg["backbone"])                         #lcnn
    print("\tJunction decoder: %s" % model_cfg["junction_decoder"])         #superpoint_decoder
    print("\tHeatmap decoder: %s" % model_cfg["heatmap_decoder"])           #pixel_shuffle
    print("\t-------------------------------------")

    #SOLD2Net网络模型的类的实例化参数
    return model

#输入图像，返回一个python字典包括"junctions" "heatmap" "descriptors"，分别为处理得到的结果,tensor变量
#过程:
# 图像传入，backbone网络的实例化对象进行处理输出特征图变量
#特征图分别传入SuperPointDecoder、heatmapDecoder、DescriptorDecoder的类实例化对象，分别得到相应的结果,tensor变量
class SOLD2Net(nn.Module):
    """ Full network for SOLD². """
    def __init__(self, model_cfg):
        super(SOLD2Net, self).__init__()
        self.name = model_cfg["model_name"]             #lcnn_simple
        self.cfg = model_cfg

        # List supported network options
        self.supported_backbone = ["lcnn", "superpoint", "superpoint_4"]
        # backbone网络模型的类的实例化对象，feat_channel通道数
        #Hourglass网络类对象 feat_channel=256 大小为原图的1/4
        #SuperPoint网络类对象 feat_channel=128 大小为原图的1/8
        self.backbone_net, self.feat_channel = self.get_backbone()

        # List supported junction decoder options
        self.supported_junction_decoder = ["superpoint_decoder"]
        #Hourglass backbone输入到SuperPoint解码器的类对象 输出通道为65 大小为原图的1/8
        self.junction_decoder = self.get_junction_decoder()

        # List supported heatmap decoder options
        self.supported_heatmap_decoder = ["pixel_shuffle",
                                          "pixel_shuffle_single",
                                          "MultiScaleAttention"]
        #Hourglass backbone输入，pixel_shuffle解码方式PixelShuffleDecoder解码器的类对象 输出通道为2 大小与原图一致
        self.heatmap_decoder = self.get_heatmap_decoder()

        self.angle_decoder = angle_decoder()

        self.heatmap_head = heatmap_head()

        # List supported descriptor decoder options
        if "descriptor_decoder" in self.cfg:
            self.supported_descriptor_decoder = ["superpoint_descriptor"]
        #Hourglass backbone输入，SuperpointDescriptor解码器的类对象 输出通道为128 大小为原图1/4
            self.descriptor_decoder = self.get_descriptor_decoder()

        # Initialize the model weights
        #apply(fn)将函数fn递归地应用到模块自身以及该模块的每一个子模块，该方法通常用来初始化一个模型中的参数
        self.apply(weight_init)

    # 继承的父类nn.Moudle使用类+参数的语法直接调用__call__(self)函数，__call__(self)函数中会返回forward函数的结果
    # 因此，forward函数可以直接通过类名被调用，而不用实例化对象
    def forward(self, input_images):
        # The backbone
        features = self.backbone_net(input_images)

        # junction decoder
        junctions = self.junction_decoder(features)

        # heatmap decoder
        heatmap_backbone = self.heatmap_decoder(input_images)

        # angle branch
        angle = self.angle_decoder(heatmap_backbone)

        angle = angle * 31.5 - 10.

        heatmaps = self.heatmap_head(heatmap_backbone, angle.clone().detach())

        outputs = {"junctions": junctions, "heatmap": heatmaps, "angle": angle}

        # Descriptor decoder
        if "descriptor_decoder" in self.cfg:
            outputs["descriptors"] = self.descriptor_decoder(features)

        return outputs

    def get_backbone(self):
        """ Retrieve the backbone encoder network. """
        #self.cfg["backbone"]=lcnn  self.supported_backbone=["lcnn", "superpoint"]
        if not self.cfg["backbone"] in self.supported_backbone:
            raise ValueError(
                "[Error] The backbone selection is not supported.")
        # lcnn backbone (stacked hourglass)
        if self.cfg["backbone"] == "lcnn":
            backbone_cfg = self.cfg["backbone_cfg"]
            #**backbone_cfg 将字典解开成独立的元素作为形参
            #1 4 2 1 5
            backbone = HourglassBackbone(**backbone_cfg)
            feat_channel = 256

        elif self.cfg["backbone"] == "superpoint":
            backbone_cfg = self.cfg["backbone_cfg"]
            backbone = SuperpointBackbone()
            feat_channel = 128

        elif self.cfg["backbone"]=="superpoint_4":
            backbone_cfg = self.cfg["backbone_cfg"]
            backbone = SuperpointBackbone_4()
            feat_channel = 128

        else:
            raise ValueError(
                "[Error] The backbone selection is not supported.")
        #backbone网络模型的类的实例化对象，feat_channel通道数
        return backbone, feat_channel

    def get_junction_decoder(self):
        """ Get the junction decoder. """
        if (not self.cfg["junction_decoder"]
            in self.supported_junction_decoder):
            raise ValueError(
                "[Error] The junction decoder selection is not supported.")

        # superpoint decoder
        if self.cfg["junction_decoder"] == "superpoint_decoder":
            decoder = SuperpointDecoder(self.feat_channel,          #256
                                        self.cfg["backbone"])       #lcnn
        else:
            raise ValueError(
                "[Error] The junction decoder selection is not supported.")

        return decoder

    def get_heatmap_decoder(self):
        """ Get the heatmap decoder. """
        if not self.cfg["heatmap_decoder"] in self.supported_heatmap_decoder:
            raise ValueError(
                "[Error] The heatmap decoder selection is not supported.")

        # Pixel_shuffle decoder
        if self.cfg["heatmap_decoder"] == "pixel_shuffle":
            if self.cfg["backbone"] == "lcnn":
                decoder = PixelShuffleDecoder(self.feat_channel,    #256
                                              num_upsample=2)
            elif self.cfg["backbone"] == "superpoint":
                decoder = PixelShuffleDecoder(self.feat_channel,
                                              num_upsample=3)
            elif self.cfg["backbone"] == "superpoint_4":
                decoder = PixelShuffleDecoder(self.feat_channel,
                                              num_upsample=2)
            else:
                raise ValueError("[Error] Unknown backbone option.")
        # Pixel_shuffle decoder with single channel output
        elif self.cfg["heatmap_decoder"] == "pixel_shuffle_single":
            if self.cfg["backbone"] == "lcnn":
                decoder = PixelShuffleDecoder(
                    self.feat_channel, num_upsample=2, output_channel=1)
            elif self.cfg["backbone"] == "superpoint":
                decoder = PixelShuffleDecoder(
                    self.feat_channel, num_upsample=3, output_channel=1)
            else:
                raise ValueError("[Error] Unknown backbone option.")
        elif self.cfg["heatmap_decoder"] == "MultiScaleAttention":
            if self.cfg["backbone"] == "lcnn":
                decoder = MultiScaleAttention()
            elif self.cfg["backbone"] == "superpoint":
                decoder = MultiScaleAttention()
            else:
                raise ValueError("[Error] Unknown backbone option.")
        else:
            raise ValueError(
                "[Error] The heatmap decoder selection is not supported.")

        return decoder

    def get_descriptor_decoder(self):
        """ Get the descriptor decoder. """
        if (not self.cfg["descriptor_decoder"]
            in self.supported_descriptor_decoder):
            raise ValueError(
                "[Error] The descriptor decoder selection is not supported.")

        # SuperPoint descriptor
        if self.cfg["descriptor_decoder"] == "superpoint_descriptor":
            decoder = SuperpointDescriptor(self.feat_channel)               #256
        else:
            raise ValueError(
                "[Error] The descriptor decoder selection is not supported.")

        return decoder


def weight_init(m):
    """ Weight initialization function. """
    # Conv2D
    #判断对象是否是一个已知的类型
    if isinstance(m, nn.Conv2d):
        #正态分布初始化
        #xavier保证网络输入与输出的方差相同，即输入的方差不要过小(sigmoid输出为线性)，输入的方差过大(sigmoid输出的梯度消失)
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    # Batchnorm
    elif isinstance(m, nn.BatchNorm2d):
        #初始化参数值符合正态分布
        init.normal_(m.weight.data, mean=1, std=0.02)
        #初始化为常量
        init.constant_(m.bias.data, 0)
    # Linear
    elif isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight.data)
        if m.bias is not None:
            init.normal_(m.bias.data)
    else:
        pass
